package com.lfy.ai.qwen;
import com.alibaba.dashscope.aigc.generation.Generation;
import com.alibaba.dashscope.aigc.generation.GenerationParam;
import com.alibaba.dashscope.aigc.generation.GenerationResult;
import com.alibaba.dashscope.common.Message;
import com.alibaba.dashscope.common.ResultCallback;
import com.alibaba.dashscope.common.Role;
import com.alibaba.dashscope.exception.ApiException;
import com.alibaba.dashscope.exception.InputRequiredException;
import com.alibaba.dashscope.exception.NoApiKeyException;
import com.alibaba.dashscope.tokenizers.*;
import com.alibaba.dashscope.utils.Constants;
import com.alibaba.dashscope.utils.JsonUtils;
import com.lfy.ai.*;
import com.lfy.domain.BuilderChatResult;
import com.lfy.modules.ai.cache.AiDispatchCache;
import com.lfy.modules.common.custom.constants.AIEnum;
import com.lfy.modules.common.domain.BaseChatVO;
import com.lfy.modules.common.domain.AiChatVO;
import com.lfy.modules.common.config.sys.DictionariesKey;
import com.lfy.modules.common.config.sys.SysDictionaries;
import com.lfy.websocket.cache.ChannelCache;
import io.netty.channel.Channel;
import io.reactivex.Flowable;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.context.annotation.Import;
import org.springframework.stereotype.Component;
import org.springframework.util.CollectionUtils;

import javax.annotation.Resource;
import java.util.*;
import java.util.concurrent.Semaphore;
/**
 * 通义千问
 *
 * @author <a href="https://www.lvfeiy.top">lvfeiyang</a>
 * @date 2024/6/23 2:16
 */
@Slf4j
@Component
@ConditionalOnProperty(name = "ai.qwen.enable", havingValue = "true", matchIfMissing = false)
@Import({QwenGroupDialogue.class})
public class QwenChat extends AiChatStandard {
    static {
        String key = SysDictionaries.getValue(DictionariesKey.AI_QWEN_API_KEY);
        if(StringUtils.isBlank(key)){
            key = Key.get(AIEnum.QW_KEY);
            if (StringUtils.isBlank(key)) {
                throw new RuntimeException("No QWEN API key found in environment variable.");
            }
        }
        Constants.apiKey = key;
    }
    {
        log.info("Blog-web AI 通义千问模型启动中...");
    }
    @Resource
    private AiDispatchCache aiDispatchCache;

    private static void handleGenerationResult(GenerationResult message, StringBuilder fullContent) {
        String content = message.getOutput().getChoices().get(0).getMessage().getContent();

        fullContent.append(content);

    }

    public static void streamCallWithMessage(Generation gen, Message userMsg)
            throws NoApiKeyException, ApiException, InputRequiredException {
        String qwenType = AIEnum.QWENType.turbo.getQwenType();
        GenerationParam param = buildGenerationParam(userMsg,qwenType);
        Flowable<GenerationResult> result = gen.streamCall(param);
        StringBuilder fullContent = new StringBuilder();

        result.blockingForEach(message -> handleGenerationResult(message, fullContent));

        log.info("Full content: \n{}", fullContent.toString());
    }

    private  void streamCallWithCallback(Generation gen,GenerationParam param,String clientId,Integer groupId,String url,Integer modelType)
            throws NoApiKeyException, ApiException, InputRequiredException, InterruptedException {
        Semaphore semaphore = new Semaphore(0);
        StringBuilder fullContent = new StringBuilder();

        gen.streamCall(param, new ResultCallback<GenerationResult>() {
            @Override
            public void onEvent(GenerationResult message) {
                handleGenerationResult(message, fullContent);
                String content = message.getOutput().getChoices().get(0).getMessage().getContent();
                Channel channel = ChannelCache.getChannel(clientId);
                if (channel!= null && channel.isActive()) {
                    AiChatVO aiChatVO = new AiChatVO();
                    aiChatVO.setContent(content);
                    aiChatVO.setModelId(AIEnum.AiModel.QwenAi.getCode());
                    aiChatVO.setGroupId(groupId);
                    BaseChatVO<AiChatVO> chattingVO = BaseChatVO.getChattingVO(aiChatVO);
                    channel.writeAndFlush(BuilderChatResult.getAIResponseVO(chattingVO));
                }
            }

            @Override
            public void onError(Exception err) {
                log.error("通义千问 api交互错误Exception occurred: {}", err.getMessage());
                Channel channel = ChannelCache.getChannel(clientId);
                if (channel!= null && channel.isActive()) {
                    channel.writeAndFlush(BuilderChatResult.getAiExceptionResponseVO(err.getMessage(),groupId,AIEnum.AiModel.QwenAi.getCode()));
                }
                semaphore.release();
            }

            @Override
            public void onComplete() {
                Channel channel = ChannelCache.getChannel(clientId);
                if (channel!= null && channel.isActive()) {
                    channel.writeAndFlush(BuilderChatResult.getAiEndResponseVO(groupId,AIEnum.AiModel.QwenAi.getCode()));
                }
                semaphore.release();
            }
        });

        semaphore.acquire();
        try {
            Message userMsg = Message.builder().role(Role.ASSISTANT.getValue()).content(fullContent.toString()).build();
            AiGroupNorm aiGroup = AiGroupManager.getAiGroup(AIEnum.AiModel.QwenAi.getCode());

            List<Message> messages = Optional.ofNullable(aiGroup).map(aiGroup1 -> aiGroup1.getMessages(clientId, groupId)).orElse(Collections.emptyList());
            if(CollectionUtils.isEmpty(messages)){
                log.error("QWEN in streamCallWithCallback method,messages is not exist,clientId:{},groupId:{}",clientId,groupId);
                return;
            }
            Optional.ofNullable(aiGroup).ifPresent(aiGroup1 -> aiGroup1.putMessage(clientId, groupId, userMsg));
            tokenizerAfter(modelType,messages,url);


        }catch (Exception e){
            log.error("通义千问 tokenizerAfter error:{},url:{},question:{},modelType:{}",e,url,fullContent.toString(),modelType);
        }

//        log.info("Full content: \n{}", fullContent.toString());
    }
    private static GenerationParam buildGenerationParam(Message userMsg,String model) {

        return GenerationParam.builder()
                .model(model)
                .messages(Arrays.asList(userMsg))
                .resultFormat(GenerationParam.ResultFormat.MESSAGE)
                .topP(0.8)
                .incrementalOutput(true)
                .build();
    }
    private static GenerationParam buildGenerationParam(List<Message> msgList, String model) {

        return GenerationParam.builder()
                .model(model)
                .messages(msgList)
                .resultFormat(GenerationParam.ResultFormat.MESSAGE)
                .topP(0.8)

                .incrementalOutput(true)
                .build();
    }


    /**
     * token计费
     * @return
     */
    @Override
    public  Integer tokenizer(String question,Integer aiModelType)  {
        Tokenization tokenizer = new Tokenization();
        Message userMsg = Message.builder()
                .role(Role.USER.getValue())
                .content(question)
                .build();
        GenerationParam param = GenerationParam.builder()
                .model(AIEnum.QWENType.getQWENType(aiModelType))
                .messages(Arrays.asList(userMsg))
                .build();
        TokenizationResult result = null;
        try {
            result = tokenizer.call(param);
        } catch (Exception e) {
            log.error("通义千问 计费token异常，type:{},question:{},exception:{}",AIEnum.QWENType.getQWENType(aiModelType),question,e.getMessage());
        }
        TokenizationUsage usage = result.getUsage();
        if(Objects.isNull(usage)){
            log.error("通义千问 远程api异常 计费token为空，type:{},question:{}",AIEnum.QWENType.getQWENType(aiModelType),question);
            throw new IllegalArgumentException("ali api exception");
        }
        Integer usageInputTokens = usage.getInputTokens();
        log.info("通义千问 计费token:{},type:{},question:{}",usageInputTokens,AIEnum.QWENType.getQWENType(aiModelType),question);
        return usageInputTokens;
    }
    private Integer getSystemToken(List<Message> messages,Integer aiModelType) {
        Tokenizer tokenizer = TokenizerFactory.qwen();

        try {
            if(CollectionUtils.isEmpty(messages)){
                log.error("QWEN in getSystemToken method,messages is not exist");
                return 0;
            }
            Message message = messages.get(messages.size() - 1);
            List<Integer> list = Optional.ofNullable(message).map(m -> tokenizer.encodeOrdinary(m.getContent())).orElse(Collections.emptyList());
            log.info("通义千问回答本次消耗：{}个token",list.size());
            return list.size();
        } catch (Exception e) {
            log.error("通义千问 计费token异常，type:{},messages:{},exception:{}",AIEnum.QWENType.getQWENType(aiModelType),messages,e.getMessage());
        }
        return 0;
    }
    @Override
    protected void tokenizerAfter(Integer aiModelType,List<Message> messages,String url) {
        Integer tokenizer = getSystemToken(messages, aiModelType);
        Long userId = aiDispatchCache.getUserId(url);
        aiDispatchCache.updateUserCapacity(AIEnum.AiModel.QwenAi.getCode(), aiModelType,userId,Long.valueOf(tokenizer));
        aiDispatchCache.updateUserCapacityCache(AIEnum.AiModel.QwenAi.getCode(), aiModelType,userId,Long.valueOf(tokenizer));
    }


    @Override
    protected boolean choice(Integer aiCode) {
        return AIEnum.AiModel.QwenAi.getCode().equals(aiCode);
    }

    @Override
    public void ask(String question,Integer aiModelType,String clientId,Integer groupId,String url) {
        Generation gen = new Generation();
        Message userMsg = Message.builder().role(Role.USER.getValue()).content(question).build();
        AiGroupNorm group = AiGroupManager.getAiGroup(AIEnum.AiModel.QwenAi.getCode());
        if(Objects.isNull(group)){
            log.error("QWEN in ask method,group is not exist,modelId:{},groupId:{}",aiModelType,groupId);
           throw new IllegalArgumentException("group not found");

        }
        group.putMessage(clientId,groupId,userMsg);
        List<Message> messages = group.getMessages(clientId, groupId);
        String qwenType = AIEnum.QWENType.getQWENType(aiModelType);
        GenerationParam params = buildGenerationParam(messages, qwenType);
        try {

            streamCallWithCallback(gen, params,clientId,groupId,url,aiModelType);
        } catch (Exception e) {
            log.error("An exception occurred: {}", e.getMessage());
        }

    }
}
